import os
import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize OpenAI client
client = OpenAI(
    api_key=os.getenv("LLM_API_KEY"),
    base_url=os.getenv("LLM_BASE_URL"),
)

def generate_ideas(category, constraints, num_ideas):
    yield "🔄 正在生成创意想法..."
    prompt = f"为{category}生成{num_ideas}个创意且实用的想法。考虑这些约束条件：{constraints}。请用中文回答。"
    try:
        response = client.chat.completions.create(
            model=os.getenv("LLM_MODEL"),
            messages=[{"role": "user", "content": prompt}],
            max_tokens=500,
        )
        ideas = response.choices[0].message.content.strip()
        yield ideas
    except Exception as e:
        yield f"❌ Error: {str(e)}"

# Gradio interface
iface = gr.Interface(
    fn=generate_ideas,
    inputs=[
        gr.Dropdown(
            label="类别/项目 (Category/Project)",
            choices=["应用开发", "营销活动", "个人项目", "产品设计", "教育培训", "其他"],
            info="选择预设类别或输入自定义"
        ),
        gr.Textbox(label="约束条件 (Constraints)", placeholder="如预算有限、时间紧迫、技术栈等"),
        gr.Slider(label="想法数量 (Number of Ideas)", minimum=1, maximum=10, step=1, value=5)
    ],
    outputs=gr.Markdown(label="生成的想法 (Generated Ideas)"),
    title="AI 创意提示生成器",
    description="输入类别和约束条件，生成创意想法。",
    submit_btn="生成想法",
    clear_btn="清除",
    theme=gr.themes.Soft(),
    examples=[
        ["应用开发", "预算有限，使用免费工具", 3],
        ["营销活动", "面向年轻群体", 5],
        ["个人项目", "学习编程", 4]
    ]
)

if __name__ == "__main__":
    iface.launch()
